package com.wdf.rule.drools;

import javax.tools.*;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 类编译器
 *
 * @author moon
 * @date 2023-02-15 20:09
 * @since 1.8
 */
public class CustomClassCompiler {

    private JavaCompiler javac;
    private DynamicClassLoader classLoader;
    private Iterable<String> options;
    boolean ignoreWarnings = false;

    /**
     * 缓存待编译的源码
     */
    private Map<String, CustomSourceCode> sourceCodes = new ConcurrentHashMap<>(16);

    /**
     * 缓存生成的类
     */
    private Map<String, Class<?>> classMap = new ConcurrentHashMap<>(16);

    /**
     * 单例编译器
     */
    private static volatile CustomClassCompiler compiler;

    /**
     * 获取实例
     * @return
     */
    public static CustomClassCompiler newInstance(ClassLoader parent) {
        if ( null==compiler ){
            synchronized (CustomClassCompiler.class){
                if (null==compiler){
                    compiler = new CustomClassCompiler();
                    if (null!=parent){
                        compiler.useParentClassLoader(parent);
                    }
                }
            }
        }
        return compiler;
    }

    /**
     * 默认类加载器 (私有化构造)
     */
    private CustomClassCompiler() {
        this.javac = ToolProvider.getSystemJavaCompiler();
        this.classLoader = new DynamicClassLoader(ClassLoader.getSystemClassLoader());
    }

    /**
     * 使用父类加载器
     * @param parent
     * @return
     */
    private CustomClassCompiler useParentClassLoader(ClassLoader parent) {
        this.classLoader = new DynamicClassLoader(parent);
        return this;
    }

    /**
     * @return the class loader used internally by the compiler
     */
    public ClassLoader getClassloader() {
        return classLoader;
    }

    /**
     * Options used by the compiler, e.g. '-Xlint:unchecked'.
     *
     * @param options
     * @return
     */
    public void useOptions(String... options) {
        this.options = Arrays.asList(options);
    }

    /**
     * 忽略警告信息
     */
    public void ignoreWarnings() {
        ignoreWarnings = true;
    }

    /**
     * 向编译器添加源码
     * @param className
     * @param sourceCode
     * @return
     * @throws Exception
     */
    public CustomClassCompiler addSource(String className, String sourceCode) {
        sourceCodes.put(className, new CustomSourceCode(className, sourceCode));
        return this;
    }

    /**
     * 编译源码
     *
     * @param classNames
     * @return
     */
    public boolean compile(String ... classNames) {
        try {
            compileByNames(Arrays.asList(classNames));
            return true;
        } catch (Exception e) {
            System.out.println("Compile Exception:" + e.getMessage());
            return false;
        }
    }

    /**
     * 获取类
     * @param className
     * @return
     */
    public Class<?> getClassByName(String className){
        return classMap.get(className);
    }

    /**
     * 编译源码
     *
     * @return Map containing instances of all compiled classes
     * @throws Exception
     */
    private void compileByNames(List<String> classNames) throws Exception {

        if (sourceCodes.size() == 0) {
            throw new RuntimeException("No source code to compile");
        }

        //获取待编译类源码
        Collection<CustomSourceCode> compilationUnits;
        Set<String> keyTemp = null;
        if (null != classNames && classNames.size() > 0){
            compilationUnits = new ArrayList<>(classNames.size());
            keyTemp = new HashSet<>(classNames.size());
            for (String key:classNames){
                if (sourceCodes.containsKey(key)){
                    keyTemp.add(key);
                    compilationUnits.add(sourceCodes.get(key));
                }
            }
        } else {
            keyTemp = sourceCodes.keySet();
            compilationUnits = sourceCodes.values();
        }

        //检测源码是否全部存在
        if (keyTemp.size() < classNames.size()){
            throw new RuntimeException("Some source code not exist");
        }

        //定义警告和错误信息输出集合
        DiagnosticCollector<JavaFileObject> collector = new DiagnosticCollector<>();
        CustomJavaFileManager fileManager = new CustomJavaFileManager(javac.getStandardFileManager(null, null, StandardCharsets.UTF_8), classLoader);
        JavaCompiler.CompilationTask task = javac.getTask(null, fileManager, collector, options, null, compilationUnits);
        //编译
        boolean result = task.call();
        //编译结果
        if (!result || collector.getDiagnostics().size() > 0) {
            StringBuffer exceptionMsg = new StringBuffer();
            exceptionMsg.append("Unable to compile the source");
            boolean hasWarnings = false;
            boolean hasErrors = false;
            for (Diagnostic<? extends JavaFileObject> d : collector.getDiagnostics()) {
                switch (d.getKind()) {
                    case NOTE:
                    case MANDATORY_WARNING:
                    case WARNING:
                        hasWarnings = true;
                        break;
                    case OTHER:
                    case ERROR:
                        hasErrors = true;
                        break;
                    default:
                        break;
                }
                exceptionMsg.append("\n").append("[kind=").append(d.getKind());
                exceptionMsg.append(", ").append("line=").append(d.getLineNumber());
                exceptionMsg.append(", ").append("message=").append(d.getMessage(Locale.US)).append("]");
            }
            //是否忽略警告
            if ((hasWarnings && !ignoreWarnings ) || hasErrors) {
                throw new RuntimeException(exceptionMsg.toString());
            }
        }

        //遍历并缓存编译后的源码
        for (String className : keyTemp) {
            classMap.put(className, classLoader.loadClass(className));
        }
    }

}


